package com.LeeCode;

import java.util.ArrayList;
import java.util.List;

/**
 * 好叶子节点对的数量
 */

public class Code1530 {
    public static void main(String[] args) {
        Integer[] arr = {1, 2, 3, null, 4};
        TreeNode root = Utils.buildTree(arr);
        System.out.println(new Code1530().countPairs(root, 3));
    }

    private int count = 0;
    private int distanceLimit;

    public int countPairs(TreeNode root, int distance) {
        this.distanceLimit = distance;
        dfs(root);
        return count;
    }

    // 返回值：从当前节点出发，到其所有叶子节点的距离列表
    private List<Integer> dfs(TreeNode node) {
        if (node == null) {
            return new ArrayList<>();
        }

        // 叶子节点：只包含自己到父节点的距离 0
        if (node.left == null && node.right == null) {
            List<Integer> leaf = new ArrayList<>();
            leaf.add(0);  // 当前是叶子节点，离父节点的距离是 0
            return leaf;
        }

        List<Integer> leftDists = dfs(node.left);
        List<Integer> rightDists = dfs(node.right);

        // 合并左右子树的距离，并 +1 表示向上走了一层
        List<Integer> merged = new ArrayList<>();

        for (int d : leftDists) {
            if (d + 1 < distanceLimit) {
                merged.add(d + 1);
            }
        }
        for (int d : rightDists) {
            if (d + 1 < distanceLimit) {
                merged.add(d + 1);
            }
        }

        // 统计左右子树之间的叶子对
        for (int l : leftDists) {
            for (int r : rightDists) {
                if (l + r + 2 <= distanceLimit) {  // 加上当前节点到左右子节点的边
                    count++;
                }
            }
        }

        return merged;
    }
}
